# 模型训练
from model3 import Network
import torch
from torch import nn
from torchvision import transforms
from torchvision import datasets
from torch.utils.data import DataLoader
from torch import optim

model = Network()  # 创建模型对象
optimizer = optim.Adam(model.parameters())  # 优化模型中的参数
criterion = nn.CrossEntropyLoss()  # 分类问题，使用交叉熵损失误差
transform = transforms.Compose([
    transforms.Resize((28, 28)),  # 调整照片大小
    transforms.Grayscale(num_output_channels=1),  # 转换为单通道灰度图
    transforms.ToTensor(),  # 转换为张量
])
train_dataset = datasets.MNIST(root='', train=True, transform=transform, download=True)
# 设置小批量处理，每组有64个数据
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
print(len(train_loader))  # 训练集60000个数据，就分为938组
# 进入模型的迭代循环
for epoch in range(100):  # 外层循环，整个训练数据集的遍历次数
    # 内存循环进行小批量的数据读取
    for batch_idx, (data, label) in enumerate(train_loader):
        # 内层每循环一次，就会进行一次梯度下降算法。
        # 包括了5个步骤，这5个步骤是使用pytorch框架训练模型的定式:
        output = model(data)  # 1.计算神经网络的前向传播结果
        loss = criterion(output, label)  # 2.计算预测值和真实值之间的差异
        loss.backward()  # 3.使用反向传播计算梯度
        optimizer.step()  # 4.使用优化器更新参数
        optimizer.zero_grad()  # 5.将梯度清零
        if batch_idx % 100 == 0:  # 每迭代100个小批量，就打印一次模型的损失，观察训练的过程
            print(f"循环次数 {epoch + 1}/100 "
                  f"| 数据集批次{batch_idx}/{len(train_loader)} "
                  f"| 损失函数值: {loss.item():.2f}")  # 值越小，表示模型的预测越好
torch.save(model, '模型文件.pth')  # 保存模型
print("模型保存成功")